#!/bin/env python3
import ithildin as ith
import numpy as np
import matplotlib.pyplot as plt
import sys
from typing import List
from scipy.interpolate import RectBivariateSpline
from pydiffmap.diffusion_map import DiffusionMap
from pydiffmap import visualization as diff_visualization
import matplotlib

# BOFC model (phase_9010/PDL_9010).
# Dit bouwt verder op dimensiereductie_AP_veelpunten.py, maar dan voor BOFC
# (en dus in meer dimensies).

# Laadt eerst de data in
data = ith.SimData.from_stem("results_9010/PDL_9010")

# Vermijd randeffecten en initialiatie
def snij_randen_weg(data,variabele,begintijd):
    def begin(index):
        if index == 0:
            return begintijd
        return data.vars[variabele].shape[index]//8
    def eind(index):
        if index == 0:
            return data.vars[variabele].shape[index]
        if data.vars[variabele].shape[index] < 20:
            return data.vars[variabele].shape[index]
        return (7 * data.vars[variabele].shape[index])//8
    return data.vars[variabele][begin(0):eind(0),begin(1):eind(1),begin(2):eind(2),begin(3):eind(3)]

def onderzoek(aantal_punten,eps):
    begintijd = (41 * 7 + 32)//33 # wacht tot de vlakke golf voorbij is (TODO controleer dit beter,msch vergelijk met simulaties in plaats van movie?)
    var_u = snij_randen_weg(data, 'u', begintijd).flatten()
    var_v = snij_randen_weg(data, 'v', begintijd).flatten()
    var_w = snij_randen_weg(data, 'w', begintijd).flatten()
    var_s = snij_randen_weg(data, 's', begintijd).flatten()
    # Verlaag aantal punten om geheugengebruik en tijdsduur te beperken.
    # replace=True strikt genomen incorrect maar wel een stuk sneller
    np.random.seed(0)
    keuzes = np.random.choice(var_u.shape[0], aantal_punten, replace=True)
    var_u = var_u[keuzes]
    var_v = var_v[keuzes]
    var_w = var_w[keuzes]
    var_s = var_s[keuzes]
    # Plot het!
    print("scatter")
    #fig = plt.figure()
    #ax = fig.add_subplot(projection='3d')
    #ax.scatter(var_u,var_v,var_w,c=var_s,cmap=plt.hot())
    #plt.show()
    #plt.close()
#    plt.scatter(var_u,var_v)
    #    plt.title("phaseplot of steekproef of BOFC model (%s points)" % (aantal_punten,)) # todo engels
    #    plt.xlabel("u")
    #    plt.ylabel("v")
    #    plt.savefig("BOFC %s punten 0.04 eps 2 eigvectoren (scatter).png" % (aantal_punten,))
    #    plt.savefig("BOFC %s punten 0.04 eps 2 eigvectoren (scatter).pdf" % (aantal_punten,))
    #    plt.close()
    # Maak nu diffusion maps
    faseruimte = np.column_stack((var_u,var_v,var_w,var_s))
    np.random.seed(100)
    dmap = DiffusionMap.from_sklearn(epsilon=eps, n_evecs=4) #0.01)
    np.random.seed(200)
    print("fitting")
    dmap.fit(faseruimte)
    print("data")
    fig = diff_visualization.data_plot(dmap,show=False)
    fig.gca().set_title("Data coloured with first DC (BOFC model, %s points)" % (aantal_punten,)) # todo engels
    plt.xlabel("u")
    plt.ylabel("v")
    fig.savefig("BOFC %s punten 0.04 eps 2 eigvectoren (data, uv).png" % (aantal_punten,))
    fig.savefig("BOFC %s punten 0.04 eps 2 eigvectoren (data, uv).pdf" % (aantal_punten,))
    plt.close()
    print("embedding")
    scatter_kwargs = { 'c': dmap.dmap[:,3], 'cmap': 'viridis' }

    fig = diff_visualization.embedding_plot(dmap,show=True,dim=3,scatter_kwargs=scatter_kwargs)
    # Why 'viridis'?  No particular reason, it's copied from an example of pydiffmap.
    fig.gca().set_title("Embedding given by first three DCs, coloured by fourth (BOFC model, %s points)" % (aantal_punten,)) # todo engels
    fig.savefig("BOFC %s punten %s eps 3+1 eigenvectors (embedding, uvw, s).png" % (aantal_punten,eps))
    fig.savefig("BOFC %s punten %s eps 3+1 eigenvectors (embedding, uvw, s).pdf" % (aantal_punten,eps))
    plt.close()
    return dmap

# Zoals ‘Embedding given by first three DCs, coloured by fourth (BOFC model, %s points)’,
# maar plot ook de punten die niet gebruikt worden om de diffusion map op te stellen.
# Hopelijk maakt dat de figuur wat duidelijker ...
def plot_veel_punten(data,aantal_punten,dmap):
    begintijd = (41 * 7 + 32)//33
    # copied from 'onderzoek'
    var_u = snij_randen_weg(data, 'u', begintijd).flatten()
    var_v = snij_randen_weg(data, 'v', begintijd).flatten()
    var_w = snij_randen_weg(data, 'w', begintijd).flatten()
    var_s = snij_randen_weg(data, 's', begintijd).flatten()
    np.random.seed(0) # determinisme (TODO)
    keuzes = np.random.choice(var_u.shape[0], aantal_punten, replace=True)
    var_u = var_u[keuzes]
    var_v = var_v[keuzes]
    var_w = var_w[keuzes]
    var_s = var_s[keuzes]
    print("gesnoeid")
    
    # bereken de diffusiecoordinaten # TODO waarom vind Emacs+Python de o-trema niet goed? 
    var_phi1234 = dmap.transform(np.column_stack((var_u,var_v,var_w,var_s)))
    print("getransformeerd")
    
    # plot het!
    fig = plt.figure(figsize=(6,6))
    ax = fig.add_subplot(111,projection='3d')
    ax.scatter(var_phi1234[:,0],var_phi1234[:,1],var_phi1234[:,2],var_phi1234[:,3],c=var_phi1234[:,3],cmap='viridis')
    ax.set_title("Embedding given by first three DCs, coloured by fourth (BOFC model)") # todo engels
    ax.set_xlabel(r'$\psi_1$')
    ax.set_ylabel(r'$\psi_2$')
    ax.set_zlabel(r'$\psi_3$')
    plt.axis('tight') # ? copied from pydiffmap
    fig.savefig("BOFC 3+1 eigenvectors, extended embedding (uvw, s).png")
    #fig.savefig("BOFC 3+1 eigenvectors, extended embedding (uvw, s).pdf")
    plt.show()
    plt.close()
    

matplotlib.use("TkAgg")
# todo np.random.seed(...) is precies niet voldoende voor determinisme
#dmap = onderzoek(1000,0.04)
#dmap = onderzoek(2000,0.01)
#dmap = onderzoek(3000,0.01)
#dmap = onderzoek(4000,0.01)
dmap = onderzoek(5000,0.01)
# 6000 is te veel voor deze laptop

# Duurt lang! Bij 10000 is er trouwens al veel structuur zichtbaar
plot_veel_punten(data,300000,dmap)
# De meeste pHet is precies één ‘hoofdreeks’

#    # Voor 7000 moet ik al een tijdje wachten
#for aantal_punten in [100, 200, 500, 1000, 2000, 3000, 4000, 5000, 6000]:
#    onderzoek(aantal_punten)

# AP 6000 punten 0.04 eps 2 eigvectoren.png: twee deuken in een cirkel, een paar uitschieters
# AP 3000 punten 0.04 eps 2 eigvectoren.png: twee deukjes, een aantal uitschieters
# AP 2000 punten 0.04 eps 2 eigvectoren.png: een appel met twee ogen
# AP 1000 punten 0.04 eps 2 eigvectoren.png: een neus met een paar stippen

